import numpy as np
from PIL import Image
import gradio as gr
from pinecone import Pinecone
from collections import Counter

# Pinecone API 初始化
pinecone = Pinecone(api_key="18cb83ea-419c-4ac9-b60f-56155477b308")
index_name = "mnist-index"
index = pinecone.Index(index_name)

# 定义预测函数
def predict_digit(image):
    # 检查图像数据格式
    if isinstance(image, dict):
        if "composite" in image:
            image = image["composite"]

    image = np.array(image)

    if image.dtype != np.uint8:
        image = (image * 255).astype(np.uint8)

    try:
        # 将图像转换为灰度并调整大小
        image = Image.fromarray(image).convert('L').resize((8, 8))
    except Exception as e:
        return "图像处理错误"

    img_array = np.array(image).reshape(1, -1)
    img_array = (img_array / 255) * 16

    # 使用 Pinecone 查询
    query_data = img_array.ravel().tolist()
    results = index.query(
        vector=query_data,
        top_k=11,
        include_metadata=True
    )

    # 打印查询结果以进行调试
    print("Query Results:", results)

    # 从搜索结果中提取每个匹配项的标签
    labels = [match['metadata']['label'] for match in results['matches']]

    # 打印每个匹配结果的详细信息
    for match, label in zip(results['matches'], labels):
        print(f"id: {match['id']}, distance: {match['score']}, label: {label}")

    # 使用投票机制确定最终的分类结果
    if labels:  # 确保 labels 不为空
        final_prediction = Counter(labels).most_common(1)[0][0]
    else:
        final_prediction = "未找到匹配"

    print(final_prediction)
    return str(final_prediction)  # 返回最终预测数字

# 创建Gradio接口
with gr.Blocks() as demo:
    gr.Markdown("## 手写数字识别器")

    with gr.Row():
        drawing_input = gr.Sketchpad(label="在这里绘制手写数字")
        prediction_output = gr.Label(label="预测的数字:")

    submit_btn = gr.Button("提交")
    clear_btn = gr.Button("清除")

    def process_image(img):
        if img is None or len(img) == 0:
            return "请绘制一个数字"
        return predict_digit(img)

    def clear_canvas():
        return None

    submit_btn.click(fn=process_image, inputs=drawing_input, outputs=prediction_output)
    clear_btn.click(fn=clear_canvas, inputs=[], outputs=drawing_input)

# 启动Gradio接口
demo.launch(share=True)  # 设置share=True可以获取公共链接